x = torch.tensor([-2.7920,-2.2356,-1.8285,-1.7755, -1.7408, -1.0992,
                      -0.9828,-0.9527,-0.2743,-0.2406, -0.0918, 0.1785,
                       0.2922, 0.3702, 0.3957, 0.4205,  0.5041, 0.5499,
                       0.6154, 0.9958, 1.2607, 1.2784,  1.4159, 1.9559])
x1 = torch.tensor(x** 1, dtype=torch.float32).reshape(-1, 1)
x2 = torch.tensor(x** 2, dtype=torch.float32).reshape(-1, 1)
x3 = torch.tensor(x** 3, dtype=torch.float32).reshape(-1, 1) # 增加三次方
x_input = torch.cat([x1, x2, x3], dim=1)  # (24,3) 
y_target = 0.5*x**1 + 3.*x**2 + 0.4*x**3   # (24,1)
y_true = y_target.reshape(-1, 1)
# 构建模型
class poly_model(nn.Module):
    def __init__(self):
        super(poly_model,self).__init__()
        self.fc = nn.Linear(3,1)
    def forward(self,x):
        out = self.fc(x)
        return out
model = poly_model()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(),lr=0.001)
# 训练
for epoch in range(3000):
    out = model(x_input)
    loss = criterion(out, y_true)
    print_loss = loss.item()
    # backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print('epoch is {} loss is {}'.format(epoch, print_loss))
print(model.fc.weight.data.numpy()[0])
print(model.fc.bias.data.numpy())
# 预测结果
model.eval()
fig, ax = plt.subplots(1, 1)
ax.plot(x, y_target.squeeze().numpy(),'ob')
predict = model(x_input)
predict = predict.data.numpy()
ax.plot(x,predict,'r-', label='degree=3')
plt.xlabel('X')
plt.ylabel('y')
ax.legend(loc='upper center', frameon=False)
plt.show()
